package org.example;

import java.io.*;
import java.net.HttpURLConnection;
import java.net.MalformedURLException;
import java.net.ProtocolException;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Scanner;


public class Chat {
    private static final String HOST = "localhost";
    private static final String HOST_CLOUD = "192.144.160.232";
    private static final int PORT_11434 = 11434;
    private static final int PORT_11000 = 11000;
    private static final List<String> MODELS = Arrays.asList(
            "deepseek-r1:1.5b",
            "deepseek-r1:7b",
            "deepseek-r1:8b",
            "deepseek-r1:14b",
            "deepseek-r1:32b",
            "deepseek-r1:70b",
            "unsloth_deepseek-r1:7B",
            "thesis:8b",
            "thesis:70b",
            "thesis:14b",
            "thesis:32b"
    );
   static class Message {
        private static final String[] ROLE = {"assistant", "user", "system"};
        private String content;

        public Message(String content) {
            this.content = content;
        }
    }

    static class ModelData {
        private String model;
        private List<Message> messages;
        private String stream;

        public ModelData(String model) {
            this.model = model;
            this.messages = new ArrayList<>();
            this.stream = "true";
        }

        @Override
        public String toString() {
            StringBuilder sb = new StringBuilder();
            sb.append("{");
            sb.append("\"model\":\"" + model + "\",");
            sb.append("\"messages\":[");
            sb.append(
                    "{\"role\":\"" + messages.get(0).ROLE[2] + "\", " +
                            "\"content\":\"" + messages.get(0).content + "\"},");
            for (int i = 1; i < messages.size(); i++) {
                if (i > 1){
                    sb.append("},");
                }
                sb.append(
                        "{\"role\":\"" + messages.get(i).ROLE[i % 2] + "\"," +
                                "\"content\":\"" + messages.get(i).content + "\"");
            }
            sb.append("}]");
//            sb.append(",\\\"stream\\\":\\"\" + stream + \"");
            sb.append("}");
            return sb.toString();
        }
    }

    public static void main(String[] args) {
        Scanner scanner = new Scanner(System.in);
        //选择模型
        //System.out.print("model-->");
        //String model = scanner.nextLine();
        //选择模型:
        System.out.println("1、deepseek-r1:1.5b");
        System.out.println("2、deepseek-r1:7b");
        System.out.println("3、deepseek-r1:8b");
        System.out.println("4、deepseek-r1:14b");
        System.out.println("5、deepseek-r1:32b");
        System.out.println("6、deepseek-r1:70b");
        System.out.println("7、unsloth_deepseek-r1:7B");
        System.out.println("8、thesis:8b");
        System.out.println("9、thesis:70b");
        System.out.println("10、thesis:14b");
        System.out.println("11、thesis:32b");
        System.out.print("选择模型:");

        Integer i = scanner.nextInt();
        String model = MODELS.get(i - 1);

        //灰色
        ModelData modelData = new ModelData(model);
        //输入系统提示语
        System.out.print("system : ");
        String m = "";
        m = scanner.nextLine();
        modelData.messages.add(new Message(m));
        scanner.nextLine();
        while (true) {
            //计算总共耗时，总共输出token，以及平均每秒输出token
            long start = 0L;
            long token = 0L;
            try {
                URL url = new URL("http://" + HOST + ":" + PORT_11434 + "/api/chat");
                HttpURLConnection conn = (HttpURLConnection) url.openConnection();
                conn.setRequestMethod("POST");
                conn.setDoOutput(true);
                conn.setRequestProperty("Content-Type", "application/json; utf-8");
                conn.setRequestProperty("Accept", "application/json");
                System.out.print("user  :  ");
                m = scanner.nextLine();
                //添加用户问题
                modelData.messages.add(new Message(m));
                start =  System.currentTimeMillis();
                try (OutputStream os = conn.getOutputStream()) {
                    byte[] input = modelData.toString().getBytes(StandardCharsets.UTF_8);
                    os.write(input, 0, input.length);
                }
                int responseCode = conn.getResponseCode();
                if (responseCode == HttpURLConnection.HTTP_OK) {
                    InputStream inputStream = conn.getInputStream();
                    BufferedReader in = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8));
                    String inputLine;
                    StringBuilder sb = new StringBuilder();
                    boolean done;
                    String role;
                    while ((inputLine = in.readLine()) != null) {

                        done = inputLine.contains("\"done\":false");
                        role = inputLine.substring(inputLine.indexOf("\"message\":{\"role\":\"") + 19, inputLine.indexOf("\",\"content\":\""));

                        if (done && role.equals("assistant")) {
                            token += 1;
                            inputLine = inputLine.substring(inputLine.indexOf("\"role\":\"assistant\",\"content\":\"") + 30, inputLine.length() - 16);
                            inputLine = inputLine.replace("\\n", "\n");
                            if (inputLine.contains("\\u003c")) {
                                inputLine = inputLine.replace("\\u003c", "\u003c");
                            }
                            if (inputLine.contains("\\u003e")){
                                inputLine = inputLine.replace("\\u003e", "\u003e");
                            }
                            System.out.print(inputLine);
                            inputLine = inputLine.replace("\n", "\\n");
                            if (inputLine.contains("\u003c")) {
                                inputLine = inputLine.replace("\u003c", "\\u003c");
                                inputLine = inputLine.replace("\u003e", "\\u003e");
                            }
                            sb.append(inputLine);
                            //###

                        }
                    }
                    modelData.messages.add(new Message(sb.toString()));
                    System.out.println();
                    in.close();
                } else {
                    System.out.println(modelData);
                    System.out.println("失败" + responseCode);
                }
            } catch (ProtocolException e) {
                throw new RuntimeException(e);
            } catch (MalformedURLException e) {
                throw new RuntimeException(e);
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
            long end = System.currentTimeMillis();
            System.out.println();
            System.out.println("*****************************************");
            System.out.println("|\t总耗时：" + ((end - start)/1000) + "s");
            System.out.println("|\t总共输出token：" + token);
            System.out.println("|\t平均每秒输出token：" + token / ((end - start) / 1000));
            System.out.println("*****************************************");
            System.out.println();
        }
    }
}
